import os
import subprocess
import shutil
import re
import json
import pygraphviz as pgv
import networkx as nx
import torch
import numpy as np

from pathlib import Path


JOERNPATH = '/home/hls/joern-cli/'
task = 'cpg14'


def check_folder(path):
    if not Path(path).is_dir():
        os.makedirs(path)


def del_file(path):
    if Path(path).is_dir():
        shutil.rmtree(path)


def packcode(code1, code2):
    def code2cpp(code, name):
        cppfile = f'{name}.cpp'
        check_folder(name)
        with open(name + '/' + cppfile, 'w') as f:
            f.write(code)
    code2cpp(code1, 'code1')
    code2cpp(code2, 'code2')


def parse_code_to_dot():
    # parse source code into cpg
    code1 = 'code1'
    code2 = 'code2'
    check_folder('dot')

    def parse(code, name):
        print(f'parseing source code into {task}...')
        shell_str = "sh " + JOERNPATH + "joern-parse " + code
        subprocess.call(shell_str, shell=True)
        print(f'exporting cpg from {task} root...')

        outpath = f'dot/{name}'
        del_file(outpath)
        shell_export = "sh " + JOERNPATH + "joern-export " + \
            f"--repr {task} --out " + outpath
        subprocess.call(shell_export, shell=True)
    parse(code1, 'code1')
    parse(code2, 'code2')
# parse的dot文件分别存放在 dot/code1 dot/code2下
# 获取main_dot
def get_main_dict(graphdir):
    mainname = ""
    alldots = list(Path(graphdir).glob('*.dot'))
    r = '^digraph\s(main)\s{'
    for dot in alldots:
        path = str(dot.parent)
        name = dot.name
        #print(name)
        class_id = path.split('\\')[-1]
        t = str(name.split('-')[0])
        with dot.open() as f:
            cpg = f.readline()
        #print(cpg)
        if re.match(r, cpg):
            #print(name)
            mainname = name
    return mainname
   
def get_main_dot():
    code1 = 'dot/code1'
    code2 = 'dot/code2'
    code1main = get_main_dict(code1)
    code2main = get_main_dict(code2)
    print(code1 + '/' + code1main)
    print(code1 + '/' + code2main)
    shutil.copyfile(code1 + '/' + code1main, './code1.dot')
    shutil.copyfile(code2 + '/' + code2main, './code2.dot')

def get_graph_massage(g, labeled=True):
        """
        输入一张图，返回 edges 和 labels
        """
        num_nodes = g.number_of_nodes()
        num_edges = g.number_of_edges()
        node_labels = ["" for i in range(num_nodes)]
        edge_labels = []

        # 给节点进行重新编号
        id = 0
        new_id = {}
        # 建立node 的label和 node的id之间的映射
        for node in g.nodes():
            new_id[node] = id
            id += 1
        
        if labeled == True:
            hsh = list(g.nodes())
            for node in g.nodes().data(): #形如 ('7', {'type': 'C', 'label': '7'})
                # labels[7] = 8
                ori_id = node[0]
                now_id = new_id[ori_id]
                if node[1] :
                    #node_labels[now_id] = node[1]["label"]
                    node_labels[now_id] = node[1]['label'].split(',')[0][1:]
                    # for condition of "<operator>.lessThan"
                    if '<operator>' in node_labels[now_id]:
                        node_labels[now_id] = node_labels[now_id].split('.')[1]
                else:
                    node_labels[now_id] = "@"
        edges = []
        # edges 的每个元素形如： [7, 4]
        for edge in g.edges().data(): #形如 ('7', '4', {'valence': 2, 'id': '6'})
            edges.append([new_id[edge[0]], new_id[edge[1]]])
            edge_labels.append(edge[2]["label"])
        # 返回一个图的边矩阵, 以及顶点向量；
        # 这是json的组成形式
        nodes = list(new_id.values())
        return nodes, edges, node_labels, edge_labels
def getgraph(dotfile):
    Gtmp1 = pgv.AGraph(dotfile)
    g = nx.DiGraph(Gtmp1)
    return g
def getpair():
    g1 = getgraph('code1.dot')
    nodes1, edges1, nodes_label1, edges_label1 = get_graph_massage(g1, labeled=True)
    g2 = getgraph('code2.dot')
    nodes2, edges2, nodes_label2, edges_label2 = get_graph_massage(g2, labeled=True)
    ret = {}
    ret["graph1"] = edges1
    ret["graph2"] = edges2
    ret["nodes1"] = nodes1
    ret["nodes2"] = nodes2
    # 每个节点的文本内容
    ret["nodes_label1"] = nodes_label1
    ret["nodes_label2"] = nodes_label2
    # 每条边的属性
    ret["edges_label1"] = edges_label1
    ret["edges_label2"] = edges_label2
    with open('pair.json', 'w') as jf:
        json.dump(ret, jf)
    return ret

def transfer_to_torch(data, global_labels, device='cpu'):
    # 转换图为无向图
    # 边记为 e ， 则edges_1的尺寸为 2e * 2
    edges_1 = data["graph1"] + [[y, x] for x, y in data["graph1"]]
    edges_2 = data["graph2"] + [[y, x] for x, y in data["graph2"]]

    # # edges_1_T的尺寸为 2* 2e; 上面的行是出发顶点；下面的行是目标顶点
    edges_1 = torch.from_numpy(np.array(edges_1, dtype=np.int64).T).type(
        torch.long)  # .to(self.device)
    edges_2 = torch.from_numpy(np.array(edges_2, dtype=np.int64).T).type(
        torch.long)  # .to(self.device)

    features_1, features_2 = [], []
    for text1 in data["nodes_label1"]:
        features_1.append(global_labels[text1])
    for text2 in data["nodes_label2"]:
        features_2.append(global_labels[text2])
    features_1 = torch.LongTensor(np.array(features_1))  # .to(self.device)
    features_2 = torch.LongTensor(np.array(features_2))  # .to(self.device)
    # edge_index_1和edge_index_2 ： 无向图表示形式的所有的边  => [2e, 2]
    return {
        "edge_index_1": edges_1.to(device),
        "edge_index_2": edges_2.to(device),
        "features_1": features_1.to(device),
        "features_2": features_2.to(device),
    }

def preprocess(code1, code2, global_labels):
    packcode(code1, code2)
    parse_code_to_dot()
    get_main_dot()
    ret = getpair()
    ret = transfer_to_torch(ret, global_labels)
    return ret
    

#preprocess(code1, code2)






# if __name__ == '__main__':
